package self.zhangkang.chart.service.index

import self.zhangkang.chart.data.KLineInfo
import self.zhangkang.chart.data.RSI
import self.zhangkang.chart.data.RSIValue
import kotlin.math.abs

/**
 * @author : Android-张康
 * created on: 2023/5/9 11:11
 * description: 计算股票的 RSI
 *
 * @param N 周期
 * @param getClosePrice 获取收盘价
 * @param getRSIValue 获取 RSI 值
 */
class RSITools<T>(
    private val N: Int,
    private val getClosePrice: (T) -> Float,
    private val getRSIValue: (T) -> RSIValue
) {

    companion object {
        //RSI参数
        private const val N1 = 6
        private const val N2 = 12
        private const val N3 = 24

        fun calculationRSI(dataList: MutableList<KLineInfo>) {
            val getClosePrice: (KLineInfo) -> Float = {
                it.closePrice
            }
            val rsi6Tools = RSITools(
                N = N1,
                getClosePrice = getClosePrice,
                getRSIValue = {
                    it.rsi.rsiValue6
                })
            val rsi12Tools = RSITools(
                N = N2,
                getClosePrice = getClosePrice,
                getRSIValue = {
                    it.rsi.rsiValue12
                })
            val rsi24Tools = RSITools(
                N = N3,
                getClosePrice = getClosePrice,
                getRSIValue = {
                    it.rsi.rsiValue24
                })
            dataList.forEachIndexed { index, kLineInfo ->
                val rsi6 = rsi6Tools.calculationRSI(
                    data = dataList,
                    currentIndex = index
                )
                val rsi12 = rsi12Tools.calculationRSI(
                    data = dataList,
                    currentIndex = index
                )
                val rsi24 = rsi24Tools.calculationRSI(
                    data = dataList,
                    currentIndex = index
                )
                dataList[index] = kLineInfo.copy(
                    rsi = RSI(rsiValue6 = rsi6, rsiValue12 = rsi12, rsiValue24 = rsi24)
                )
            }
        }
    }

    /**
     * 计算RSI
     *
     * @param data 当前的数据
     * @param currentIndex 当前数据
     */
    fun calculationRSI(
        data: List<T>,
        currentIndex: Int,
    ): RSIValue {
        val diff: Float
        val preIn: Float
        val preAll: Float
        if (currentIndex == 0) {
            diff = 0F
            preIn = 0F
            preAll = 0F
        } else {
            val previousData = data[currentIndex - 1]
            diff = getClosePrice(data[currentIndex]) - getClosePrice(previousData)
            val rsiValue = getRSIValue(previousData)
            preIn = rsiValue.averageIncrease
            preAll = rsiValue.averageDecline
        }
        val averageIncrease = countSMA(maxOf(diff, 0F), N, 1, preIn)
        val averageDecline = countSMA(abs(diff), N, 1, preAll)
        return if (currentIndex >= N) {
            RSIValue(
                averageIncrease = averageIncrease,
                averageDecline = averageDecline,
                rsi = averageIncrease / averageDecline * 100
            )
        } else {
            RSIValue(
                averageIncrease = averageIncrease,
                averageDecline = averageDecline,
                rsi = Float.NaN
            )
        }
    }

    private fun countSMA(c: Float, n: Int, m: Int, sma: Float): Float {
        return (m * c + (n - m) * sma) / n
    }
}